"""
使用yolo v10 进行检测
使用摄像头输入实时处理并返回结果
"""

import gradio as gr
import torch
from ultralytics import YOLO

print(f"PyTorch 版本: {torch.__version__}")
print(f"CUDA 可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA 版本: {torch.version.cuda}")
    print(f"GPU 设备: {torch.cuda.get_device_name(0)}")


model = YOLO(f'c:/tmp/yolo_v10s.pt')

base_conf, base_iou = 0.8, 0.2

def detect(image, conf, iou):
    if image is None or type(image) == 'NoneType':
        return None

    results = model(source=image,conf=conf,iou=iou)
    annotated_frame = results[0].plot()
    return annotated_frame


css = """
footer {display: none !important;}
"""

demo = gr.Interface(
    inputs=[gr.Image(sources=["webcam"], streaming=True, width=640, height=480),gr.Slider(minimum=0, maximum=1, value=base_conf,label='target confidence show threshold'),gr.Slider(minimum=0, maximum=1, value=base_iou,label='NMS IOU show threshold')],
    outputs=[gr.Image()],
    fn=detect,
    live=True,
    title="where is my son?",
)

demo.launch(css=css,share=False)
# demo.launch(css=css, server_name="0.0.0.0")
